/**
 * @file main.cpp
 *
 * Copyright (C) 2025. Huawei Technologies Co., Ltd. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */
#include "reflection_pad1d_tiling.h"  
#include "data_utils.h"      
#include <cstring>
#include <iostream>

#ifndef ASCENDC_CPU_DEBUG
#include "acl/acl.h"
#include "aclrtlaunch_reflection_pad1d.h"
#include "tiling/platform/platform_ascendc.h"
#else
#include "tikicpulib.h"
extern "C" __global__ __aicore__ void reflection_pad1d(GM_ADDR x, GM_ADDR paddings, GM_ADDR y, ReflectionPad1dTilingData tiling);
#endif

extern void GenerateTilingData(uint8_t *tilingBuf, uint32_t coreNum);

int32_t main(int32_t argc, char *argv[])
{
    constexpr uint32_t BLOCK_DIM = 8; 
    const std::string inputTilingPath = "./input/input_tiling.bin";  // 分块配置文件路径
    const std::string inputXPath = "./input/input_x.bin";            // 输入张量文件路径
    const std::string inputPaddingsPath = "./input/paddings.bin";    // 填充参数文件路径
    const std::string outputYPath = "./output/output_y.bin";         // 输出张量文件路径

    uint8_t *tilingBuf = nullptr;
    size_t tilingSize = sizeof(ReflectionPad1dTilingData); 
#ifdef ASCENDC_CPU_DEBUG
    tilingBuf = (uint8_t *)AscendC::GmAlloc(tilingSize);
#else
    CHECK_ACL(aclInit(nullptr));
    int32_t deviceId = 0;
    CHECK_ACL(aclrtSetDevice(deviceId));
    aclrtStream stream = nullptr; 
    CHECK_ACL(aclrtCreateStream(&stream));
    CHECK_ACL(aclrtMallocHost((void **)&tilingBuf, tilingSize));
#endif
    if (!ReadFile(inputTilingPath.c_str(), tilingSize, tilingBuf, tilingSize)) {
        std::cerr << "Error: 读取分块配置文件 " << inputTilingPath << " 失败！" << std::endl;
        return -1;
    }

    GenerateTilingData(tilingBuf, BLOCK_DIM); 
    ReflectionPad1dTilingData *tilingData = reinterpret_cast<ReflectionPad1dTilingData *>(tilingBuf);
    
    uint32_t dtypeBytes = 0;
    switch (tilingData->tilingKey) {
        case static_cast<uint32_t>(TilingKey::TILING_KEY_FP32):
            dtypeBytes = FLOAT_BYTES; 
            break;
        case static_cast<uint32_t>(TilingKey::TILING_KEY_FP16):
        case static_cast<uint32_t>(TilingKey::TILING_KEY_BF16):
            dtypeBytes = FLOAT16_BYTES; 
            break;
        default:
            std::cerr << "Error: 不支持的数据类型 tilingKey=" << static_cast<uint32_t>(tilingData->tilingKey) << std::endl;
            return -1;
    }

    size_t inputXByteSize = tilingData->nSize * tilingData->cSize * tilingData->alignWSize * dtypeBytes;
    size_t paddingsByteSize = 2 * sizeof(int32_t);  
    size_t outputYByteSize = tilingData->nSize * tilingData->cSize * (tilingData->wSize + tilingData->padLeft + tilingData->padRight) * dtypeBytes;

#ifdef ASCENDC_CPU_DEBUG
    uint8_t *x = nullptr;    // 输入张量（GM内存）
    uint8_t *paddings = nullptr;  // 填充参数（GM内存）
    uint8_t *y = nullptr;    // 输出张量（GM内存）
#else
    uint8_t *xHost = nullptr, *paddingsHost = nullptr, *yHost = nullptr;  // Host内存
    uint8_t *xDevice = nullptr, *paddingsDevice = nullptr, *yDevice = nullptr;  // Device内存
#endif
    try {
#ifdef ASCENDC_CPU_DEBUG
        x = (uint8_t *)AscendC::GmAlloc(inputXByteSize);
        paddings = (uint8_t *)AscendC::GmAlloc(paddingsByteSize);
        y = (uint8_t *)AscendC::GmAlloc(outputYByteSize);
        if (!x || !paddings || !y) {
            throw std::runtime_error("GM内存分配失败！");
        }

        if (!ReadFile(inputXPath.c_str(), inputXByteSize, x, inputXByteSize) ||
            !ReadFile(inputPaddingsPath.c_str(), paddingsByteSize, paddings, paddingsByteSize)) {
            throw std::runtime_error("读取输入文件失败！");
        }

        AscendC::SetKernelMode(KernelMode::AIV_MODE);
        ICPU_RUN_KF(reflection_pad1d, BLOCK_DIM, x, paddings, y, *tilingData);

        if (!WriteFile(outputYPath.c_str(), y, outputYByteSize)) {
            throw std::runtime_error("写入输出文件失败！");
        }
#else

        CHECK_ACL(aclrtMallocHost((void **)&xHost, inputXByteSize));
        CHECK_ACL(aclrtMallocHost((void **)&paddingsHost, paddingsByteSize));
        CHECK_ACL(aclrtMallocHost((void **)&yHost, outputYByteSize));  // 存储Device返回的结果
        CHECK_ACL(aclrtMalloc((void **)&xDevice, inputXByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
        CHECK_ACL(aclrtMalloc((void **)&paddingsDevice, paddingsByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
        CHECK_ACL(aclrtMalloc((void **)&yDevice, outputYByteSize, ACL_MEM_MALLOC_HUGE_FIRST));

        if (!ReadFile(inputXPath.c_str(), inputXByteSize, xHost, inputXByteSize) ||
            !ReadFile(inputPaddingsPath.c_str(), paddingsByteSize, paddingsHost, paddingsByteSize)) {
            throw std::runtime_error("读取输入文件失败！");
        }

        CHECK_ACL(aclrtMemcpy(xDevice, inputXByteSize, xHost, inputXByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
        CHECK_ACL(aclrtMemcpy(paddingsDevice, paddingsByteSize, paddingsHost, paddingsByteSize, ACL_MEMCPY_HOST_TO_DEVICE));

        ACLRT_LAUNCH_KERNEL(reflection_pad1d)(BLOCK_DIM, stream, xDevice, paddingsDevice, yDevice, tilingData);

        CHECK_ACL(aclrtSynchronizeStream(stream));

        CHECK_ACL(aclrtMemcpy(yHost, outputYByteSize, yDevice, outputYByteSize, ACL_MEMCPY_DEVICE_TO_HOST));

        if (!WriteFile(outputYPath.c_str(), yHost, outputYByteSize)) {
            throw std::runtime_error("写入输出文件失败！");
        }
#endif
        std::cout << "reflection_pad1d 执行成功！" << std::endl;
        std::cout << "输入张量尺寸（对齐后）：N=" << tilingData->nSize 
                  << ", C=" << tilingData->cSize << ", W(align)=" << tilingData->alignWSize << std::endl;
        std::cout << "输出张量尺寸（对齐后）：W(out,align)=" << tilingData->alignOutWSize << std::endl;
        std::cout << "填充参数：left=" << tilingData->padLeft << ", right=" << tilingData->padRight << std::endl;
        std::cout << "输出文件：" << outputYPath << std::endl;

    } catch (const std::exception &e) {
        std::cerr << "Error: " << e.what() << std::endl;
        return -1;
    }
#ifdef ASCENDC_CPU_DEBUG
    AscendC::GmFree(x);
    AscendC::GmFree(paddings);
    AscendC::GmFree(y);
    AscendC::GmFree(tilingBuf);
#else
    CHECK_ACL(aclrtFree(xDevice));
    CHECK_ACL(aclrtFree(paddingsDevice));
    CHECK_ACL(aclrtFree(yDevice));
    CHECK_ACL(aclrtFreeHost(xHost));
    CHECK_ACL(aclrtFreeHost(paddingsHost));
    CHECK_ACL(aclrtFreeHost(yHost));
    CHECK_ACL(aclrtFreeHost(tilingBuf));
    CHECK_ACL(aclrtDestroyStream(stream));
    CHECK_ACL(aclrtResetDevice(deviceId));
    CHECK_ACL(aclFinalize());
#endif
    return 0;
}